% main_fatigue
% Jones, Philippon, Venkateswaran (2021)
%

%%

clear 
clc

%%

warning('Off','all') ;

path(pathdef) ;

%% 

planner = 1 ;

e0exog   = 0 ;
Imax     = 0.2;              %% targets for mortality
BaseMort = 0.0025;           
PeakMort = 3*BaseMort;
deterministic_vac = 1 ;
stoch_prob = 77/78 ;

beta       = (0.97)^(1/52)*stoch_prob ; 
ee_T_vac   = zeros(1,1000); 
ee_T_e0    = zeros(1,1000); 
ee_T_d     = zeros(1,1000); 
if deterministic_vac; beta = (0.97)^(1/52); T_vac = 1*52 ; ee_T_vac(T_vac:end)=0.01; for ii=(T_vac+1):length(ee_T_vac); ee_T_vac(ii)=ee_T_vac(ii-1)*1.2; end ; ee_T_vac(ee_T_vac>1) = 1; end
e0         = e0exog        + (1-e0exog)*(1/2) ;
ec1        = e0exog*0.0001 + (1-e0exog)*((1-e0)/2) ;
el1        = e0exog*0.0001 + (1-e0exog)*((1-e0)/2) ;
R0         = 2.0 ;
chibar1    = 0.5 ;
kappa      = 0.15;                                    % Fraction severely sick
rho        = 0.35;                                    % Recovery rate of infected 
deltabar   = BaseMort/(1-BaseMort)*rho/kappa;         % baseline death rate 
phi        = log(1-deltabar+PeakMort/(1-PeakMort)*rho/kappa)/(Imax*kappa);
N          = 1 ;
eta        = 2 ;
rhom       = 1 ;
Deltachi   = 0.34 ;
us         = 0.5 ;
ud         = 2.5 ;
ee_i0_init = (1/100)*N ; ee_T_i = zeros(1,1000) ; ee_T_i(1) = ee_i0_init ;
psim       = 0.99 ;
deltascale = 1 ;

rhof        = 0.0 ;
omega1      = 0.0 ;
omega2      = 2 ;

rhoa       = 0.0 ;
varphi1    = 0.5 ;
varphi2    = 0.2 ;
amean      = 1 ;

if planner
    model_name = 'infections_planner_fatigue' ;
    set_params_planner
else
    model_name = 'infections' ;
    set_params
end

run_dynare

figure(1)

set(gcf, 'DefaultAxesLineWidth', 1.5);
set(gcf, 'DefaultLineLineWidth', 2);
set(gcf, 'DefaultAxesTickLabelInterpreter','latex'); 
set(gcf, 'DefaultLegendInterpreter','latex');
set(gcf, 'DefaultAxesFontSize',10);

subplot(2,2,2); hold on;
h1=plot(2:(length(s)-1),e(3:end)) ;
line([0 300],[1 1],'LineWidth',1,'Color','k') ;
title('Exposure, $e$','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ; 

subplot(2,2,1); 
yyaxis left; hold on ;
h1 = plot(0:100,C1(1:101),'-','Color',h1.Color) ;
title('Consumption (LHS), Labor (RHS)','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ; 
set(gca,'YColor','k') ;
ylim([0.8 1.05]) ;

yyaxis right; hold on ;
h1 = plot(0:100,L1(1:101),'--','Color',h1.Color) ;
title('Consumption (LHS), Labor (RHS)','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ; 
set(gca,'YColor','k') ;
ylim([0.8 1.05]) ;

subplot(2,2,3); hold on;
h1 = plot(0:(length(s)-1),100*in) ;
title('Infected, \%','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ; 

subplot(2,2,4); hold on;
yyaxis left
hold on ;
plot(0:(length(s)-1),100*s,'-','Color',h1.Color) ;
xlim([0 75]); xticks([0:25:75]) ; 
set(gca,'YColor','k') ;

yyaxis right
hold on ;
plot(0:(length(d)-1),100*d,'--','Color',h1.Color) ;
xlim([0 75]); xticks([0:25:75]) ; 
set(gca,'YColor','k') ;

title('Susceptible (LHS), Dead (RHS), \%','Interpreter','latex')

set(findall(gcf, 'Type', 'Text'),'FontWeight', 'Normal')

%% 

planner = 1 ;

e0exog   = 0 ;
Imax     = 0.2;              %% targets for mortality
BaseMort = 0.0025;           
PeakMort = 3*BaseMort;
deterministic_vac = 1 ;
stoch_prob = 77/78 ;

beta       = (0.97)^(1/52)*stoch_prob ; 
ee_T_vac   = zeros(1,1000); 
ee_T_e0    = zeros(1,1000); 
ee_T_d     = zeros(1,1000); 
if deterministic_vac; beta = (0.97)^(1/52); T_vac = 1*52 ; ee_T_vac(T_vac:end)=0.01; for ii=(T_vac+1):length(ee_T_vac); ee_T_vac(ii)=ee_T_vac(ii-1)*1.2; end ; ee_T_vac(ee_T_vac>1) = 1; end
e0         = e0exog        + (1-e0exog)*(1/2) ;
ec1        = e0exog*0.0001 + (1-e0exog)*((1-e0)/2) ;
el1        = e0exog*0.0001 + (1-e0exog)*((1-e0)/2) ;
R0         = 2.0 ;
chibar1    = 0.5 ;
kappa      = 0.15;                                    % Fraction severely sick
rho        = 0.35;                                    % Recovery rate of infected 
deltabar   = BaseMort/(1-BaseMort)*rho/kappa;         % baseline death rate 
phi        = log(1-deltabar+PeakMort/(1-PeakMort)*rho/kappa)/(Imax*kappa);
N          = 1 ;
eta        = 2 ;
rhom       = 1 ;
Deltachi   = 0.34 ;
us         = 0.5 ;
ud         = 2.5 ;
ee_i0_init = (1/100)*N ; ee_T_i = zeros(1,1000) ; ee_T_i(1) = ee_i0_init ;
psim       = 0.99 ;
deltascale = 1 ;

rhof        = 0.99 ;
omega1      = 0.5 ;
omega2      = 2 ;

rhoa       = 0.99 ;
varphi1    = 0.5 ;
varphi2    = 0.2 ;
amean      = 1 ;

if planner
    model_name = 'infections_planner_fatigue' ;
    set_params_planner
else
    model_name = 'infections' ;
    set_params
end

run_dynare


figure(1)

set(gcf, 'DefaultAxesLineWidth', 1.5);
set(gcf, 'DefaultLineLineWidth', 2);
set(gcf, 'DefaultAxesTickLabelInterpreter','latex'); 
set(gcf, 'DefaultLegendInterpreter','latex');
set(gcf, 'DefaultAxesFontSize',10);

subplot(2,2,2); hold on;
h1=plot(2:(length(s)-1),e(3:end)) ;
line([0 300],[1 1],'LineWidth',1,'Color','k') ;
title('Exposure, $e$','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ; 

subplot(2,2,1); 
yyaxis left; hold on ;
h1 = plot(0:100,C1(1:101),'-','Color',h1.Color) ;
title('Consumption (LHS), Labor (RHS)','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ; 
set(gca,'YColor','k') ;
ylim([0.7 1.05]) ;

yyaxis right; hold on ;
h1 = plot(0:100,L1(1:101),'--','Color',h1.Color) ;
title('Consumption (LHS), Labor (RHS)','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ; 
set(gca,'YColor','k') ;
ylim([0.7 1.05]) ;
h=legend('No Fatigue','With Fatigue','location','southeast');
set(h,'Interpreter','latex','FontSize',6) ;

subplot(2,2,3); hold on;
h1 = plot(0:(length(s)-1),100*in) ;
title('Infected, \%','Interpreter','latex')
xlim([0 75]); xticks([0:25:75]) ; 

subplot(2,2,4); hold on;
yyaxis left
hold on ;
plot(0:(length(s)-1),100*s,'-','Color',h1.Color) ;
xlim([0 75]); xticks([0:25:75]) ; 
set(gca,'YColor','k') ;

yyaxis right
hold on ;
plot(0:(length(d)-1),100*d,'--','Color',h1.Color) ;
xlim([0 75]); xticks([0:25:75]) ;
set(gca,'YColor','k') ;

title('Susceptible (LHS), Dead (RHS), \%','Interpreter','latex')

set(findall(gcf, 'Type', 'Text'),'FontWeight', 'Normal')

print -depsc ./output/fatigue_baseline.eps
print -dpng ./output/fatigue_baseline.png
close
